import torch.nn as nn
from src.nets.vanillaNet import vanillanet_13 as vanillaNet


class MRNet(nn.Module):
    def __init__(self, pretrained=False, backbond='vanilla', backbond_model=''):
        super().__init__()
        if backbond_model:
            self.pretrained_model = vanillaNet(pretrained=True, model_path=backbond_model)
            # 将梯度回传设置成False
            # for param in self.pretrained_model.parameters():
            #     param.requires_grad = False
        self.classify = nn.Linear(1000, 3)

    def forward(self, x):
        features = self.pretrained_model(x)
        output = self.classify(features)

        return output
